#include <fcntl.h>
#include <stdio.h>
#include <string.h>

#include <openssl/pem.h>
#include <openssl/evp.h>

#include "trace.h"

int save_private_key_dsa(const char *filename, const char *passphrase,
        DSA *dsa) {
    int ret_code = 0;

    if((NULL == filename) || (NULL == passphrase) ||
            (NULL == dsa)){
        TRACE_ERROR("Invalid input addresses\n");
        return 1;
    }

    BIO *out;
    out = BIO_new(BIO_s_file());
    if(out == NULL) {
        TRACE_ERROR("BIO_new error\n");
        return 1;
    }

    if(!BIO_write_filename(out, (char*)filename)) {
        TRACE_ERROR("BIO_write_filename error\n");
        return 2;
    }

    int len = strlen(passphrase);
    if ((len > 0) && (len <= 4)) {
        TRACE_ERROR("Passphrase too short: %d bytes\n", len);
        return 1;
    }

    if (len > 0) {
        // encrypt with provided password
        if (!PEM_write_bio_DSAPrivateKey(out, dsa, EVP_aes_256_cbc(),
                (unsigned char *)passphrase, strlen(passphrase), NULL, NULL)){
            TRACE_ERROR("Write with password error");
            return 1;
        }
    }
    else {
        // no passphrase. Use default callback to get one
        if (!PEM_write_bio_DSAPrivateKey(out, dsa, EVP_aes_256_cbc(),
                NULL, 0, NULL, NULL)){
            TRACE_ERROR("Write with password callback error\n");
            return 1;
        }

    }

    BIO_free(out);

    return ret_code;
}

int save_public_key_dsa(const char *filename, DSA *dsa) {
    int ret_code = 0;

    if((NULL == filename) || (NULL == dsa)){
        TRACE_ERROR("Invalid input addresses");
        return 1;
    }

    BIO *out;
    out = BIO_new(BIO_s_file());
    if(out == NULL) {
        TRACE_ERROR("BIO_new error\n");
        return 1;
    }

    if(!BIO_write_filename(out, (char*)filename)) {
        TRACE_ERROR("BIO_write_filename error\n");
        return 2;
    }

    if(!PEM_write_bio_DSA_PUBKEY(out, dsa)){
        TRACE_ERROR("Write public key to PEM error");
        return 1;
    }

    BIO_free(out);

    return ret_code;
}

int load_private_key_dsa(const char *filename, const char *passphrase,
        DSA **dsa) {
    int ret_code = 0;

    if((NULL == filename) || (NULL == passphrase)){
        TRACE_ERROR("Invalid input addresses\n");
        return 1;
    }

    BIO *in = NULL;
    in = BIO_new(BIO_s_file());
    if(in == NULL) {
        TRACE_ERROR("BIO_new error\n");
        return 1;
    }

	if(!BIO_read_filename(in, (char*)filename)) {
		TRACE_ERROR("BIO_read_filename error for file: %s\n", filename);
        return 2;
    }

    int len = strlen(passphrase);
    if (len > 0) {
        // decrypt private key with provided password
        if(!PEM_read_bio_DSAPrivateKey(in, dsa, NULL, (void*)passphrase)){
            TRACE_ERROR("Read with password error\n");
            return 1;
        }
    }
    else {
        // no passphrase. Use default callback to get one
        if(!PEM_read_bio_DSAPrivateKey(in, dsa, NULL, NULL)){
            TRACE_ERROR("Read with password callback error\n");
            return 1;
        }

    }

    BIO_free(in);

    return ret_code;
}

int load_public_key_dsa(const char *filename, DSA **dsa) {
    int ret_code = 0;

    if((NULL == filename) || (NULL == dsa)){
        TRACE_ERROR("Invalid input addresses\n");
        return 1;
    }

    BIO *in;
    in = BIO_new(BIO_s_file());
    if(in == NULL) {
        TRACE_ERROR("BIO_new error\n");
        return 1;
    }

    if(!BIO_read_filename(in, (char*)filename)) {
        TRACE_ERROR("BIO_read_filename error\n");
        return 2;
    }

    if(!PEM_read_bio_DSA_PUBKEY(in, dsa, NULL, NULL)){
        TRACE_ERROR("Read public key error\n");
        return 1;
    }

    return ret_code;
}

int generate_DSA_key_pair(const char* passphrase){
    int ret_code = 0;

    if(NULL == passphrase){
        TRACE_ERROR("Invalid parameter to generate_DSA_key_pair\n");
        return 1;
    }

    DSA* dsa = DSA_new();

    // generate p, q, g parameters
    DSA_generate_parameters_ex(dsa, 1024, NULL, 0, NULL, NULL, NULL);

    // generate public and private keys
    if(!DSA_generate_key(dsa)) {
        TRACE_ERROR("DSA_generate_key error\n");
        return 1;
    }

    if(save_private_key_dsa("my_prv_key.pem", passphrase, dsa)){
        TRACE_ERROR("Error saving private key\n");
        return 2;
    }
    else {
        TRACE_DEBUG("Private key saved to my_prv_key.pem\n");
    }

    if(save_public_key_dsa("my_pub_key.pem", dsa)){
        TRACE_ERROR("Error saving public key\n");
        return 3;
    }
    else{
        TRACE_DEBUG("Public key saved to my_pub_key.pem\n");
    }

    return ret_code;
}
